from typing import List, Optional, Any

from sqlalchemy.orm import Session
from fastapi.encoders import jsonable_encoder
from fastapi_pagination.ext.sqlalchemy import paginate
from fastapi_pagination import Page

from app.utils import ModelType
from app.crud.base import CRUDBase
from app.models import Group, User
from app.models.link import user_group_association
from app.schemas import GroupCreate, GroupUpdate
from app.utils import apply_changes


class CRUDDemo(CRUDBase[Group, GroupCreate, GroupUpdate]):
    def create_link_user_group(
            self, db: Session, *, obj_in: GroupCreate, user: Optional[User]
    ) -> Group:
        obj_in_data = jsonable_encoder(obj_in)
        obj_in_data.pop("users")
        db_obj = self.model(**obj_in_data)  # type: ignore
        apply_changes(db, db_obj)

        # 根据 user_ids 添加用户到组中
        for user_id in obj_in.users:
            association = user_group_association.insert().values(user_id=user_id, group_id=db_obj.id)
            db.execute(association)

        db.commit()
        return db_obj

    def get_paginate_link_user_group(
            self, db: Session, *, user: User
    ) -> Page:
        groups_for_user = (
            db.query(Group)
            .join(user_group_association, Group.id == user_group_association.c.group_id)
            .filter(user_group_association.c.user_id == user.id)
        )
        return paginate(groups_for_user)

    def update_link_user_group(
            self, db: Session, *, db_obj: ModelType, obj_in: GroupUpdate
    ) -> Group:
        """
        更新 group 并更新关联关系
        """
        # 更新 group 原表的内容
        obj_in_data = jsonable_encoder(obj_in)
        new_users = set(obj_in_data.pop("users"))
        self.update(db, db_obj=db_obj, obj_in=obj_in)

        # 找出 group 组中所有的 user
        stmt = db.query(User).join(user_group_association).where(user_group_association.c.group_id == db_obj.id)
        users = db.execute(stmt).scalars().all()
        old_users = set([str(user.id) for user in users])

        # 需要添加的关联关系
        new_users_to_add = new_users.difference(old_users)
        for user_id in new_users_to_add:
            db.execute(user_group_association.insert().values(user_id=user_id, group_id=db_obj.id))

        # 需要删除的关联关系
        old_users_to_remove = old_users.difference(new_users)
        for user_id in old_users_to_remove:
            db.execute(user_group_association.delete().where(user_group_association.c.user_id == user_id,
                                                             user_group_association.c.group_id == db_obj.id))

        db.commit()
        return db_obj


group = CRUDDemo(Group)
